{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cdfcc8c6",
   "metadata": {},
   "source": [
    "# 多分类通用模型，使用softmax计算损失函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f222a2d7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "使用设备: cpu\n",
      "已设置中文字体支持（Windows系统）\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.datasets import ImageFolder\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "from PIL import Image\n",
    "import platform\n",
    "from utils import *\n",
    "import cv2\n",
    "\n",
    "# 检查设备\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"使用设备: {device}\")\n",
    "set_chinese_font()\n",
    "# visualize_samples()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "67d946ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = r\"D:\\desktop\\class\\DataSet\\three\"\n",
    "\n",
    "# 简化版数据变换（去掉数据增强，只保留基本变换）\n",
    "simple_train_transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.Grayscale(num_output_channels=1),  # 灰度图\n",
    "        transforms.Resize((64, 64)),  # 调整大小\n",
    "        transforms.ToTensor(),  # 转换为Tensor\n",
    "        transforms.Normalize(mean=[0.5], std=[0.5])  # 单通道的标准化参数\n",
    "    ]\n",
    ")\n",
    "\n",
    "simple_test_transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.Grayscale(num_output_channels=1),  # 灰度图\n",
    "        transforms.Resize((64, 64)),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.5], std=[0.5])  # 单通道的标准化参数\n",
    "    ]\n",
    ")\n",
    "\n",
    "# 使用简化变换重新创建数据集\n",
    "train_dataset = ImageFolder(\n",
    "    os.path.join(data_path, \"train\"), transform=simple_train_transform\n",
    ")\n",
    "\n",
    "test_dataset = ImageFolder(\n",
    "    os.path.join(data_path, \"test\"),\n",
    "    transform=simple_test_transform\n",
    "\n",
    ")\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=600, shuffle=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=50, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e07746b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "类别列表: ['butterflys', 'cats', 'dogs']\n",
      "类别到索引的映射: {'butterflys': 0, 'cats': 1, 'dogs': 2}\n",
      "数据集大小: 900\n",
      "图片形状: torch.Size([1, 64, 64])\n",
      "标签: 0 -> butterflys\n",
      "图片形状: torch.Size([1, 64, 64])\n",
      "标签: 0 -> butterflys\n"
     ]
    }
   ],
   "source": [
    "print(\"类别列表:\", train_dataset.classes)\n",
    "print(\"类别到索引的映射:\", train_dataset.class_to_idx)\n",
    "print(\"数据集大小:\", len(train_dataset))\n",
    "\n",
    "# 查看第一个样本\n",
    "sample, label = train_dataset[0]\n",
    "print(\"图片形状:\", sample.shape)\n",
    "print(\"标签:\", label, \"->\", train_dataset.classes[label])\n",
    "\n",
    "sample, label = train_dataset[145]\n",
    "print(\"图片形状:\", sample.shape)\n",
    "print(\"标签:\", label, \"->\", train_dataset.classes[label])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c79033e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# show_transformed_images(8,train_dataset,simple_train_transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "002f5500",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleCatDogCNN(nn.Module):\n",
    "    def __init__(self, num_classes=3):\n",
    "        super(SimpleCatDogCNN, self).__init__()\n",
    "        self.network = nn.Sequential(\n",
    "            nn.Conv2d(1, 2, kernel_size=3, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            # 以下激活函数可以替换比较一下\n",
    "            # 'ReLU': nn.ReLU(inplace=True),\n",
    "            # 'LeakyReLU': nn.LeakyReLU(0.01, inplace=True),\n",
    "            # 'PReLU': nn.PReLU(num_parameters=2),\n",
    "            # 'Mish': nn.Mish(inplace=True),\n",
    "            # 'Tanh': nn.Tanh()\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(2 * 64 * 64, num_classes) # 使用softmax作为损失函数，因此类数为2\n",
    "        )\n",
    " \n",
    "    def forward(self, x):\n",
    "        return self.network(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05d4e606",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SimpleCatDogCNN(num_classes=2).to(device)\n",
    "print(\"模型结构:\")\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ba2682e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, train_loader, test_loader, criterion, optimizer, num_epochs=50):\n",
    "    train_losses = []\n",
    "    test_accuracies = []\n",
    "    best_test_acc = 0.0\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        # 当前情况下（batchsize = 600，只有一个batch），每一次epoch其实只执行一次梯度下降\n",
    "        # 训练阶段\n",
    "        model.train()\n",
    "        total_loss = 0.0\n",
    "        total_correct = 0\n",
    "        total_samples = 0\n",
    "        \n",
    "        for inputs, labels in train_loader:\n",
    "            # 循环一次，batch_size=600\n",
    "            # inputs: [batch_size, 1, 64, 64], labels: [batch_size], batch_size=600\n",
    "            inputs, labels = inputs.to(device), labels.to(device)\n",
    "            \n",
    "            # 前向传播 + 反向传播\n",
    "            optimizer.zero_grad()   # 清零梯度\n",
    "            outputs = model(inputs) # 前向传播\n",
    "            loss = criterion(outputs, labels)   # 计算损失\n",
    "            loss.backward() # 反向传播\n",
    "            optimizer.step()# 更新参数\n",
    "            \n",
    "            # 统计 - 修正这里\n",
    "            batch_size = inputs.size(0)  # 当前batch的大小（600）\n",
    "            total_loss += loss.item() * batch_size  # 乘以batch_size得到总损失\n",
    "            _, preds = torch.max(outputs, 1)\n",
    "            total_correct += (preds == labels).sum().item()\n",
    "            total_samples += batch_size\n",
    "\n",
    "        # 计算平均损失和准确率 - 修正这里\n",
    "        avg_loss = total_loss / total_samples  # 除以总样本数得到平均损失\n",
    "        train_acc = total_correct / total_samples\n",
    "        \n",
    "        train_losses.append(avg_loss)\n",
    "\n",
    "        # 验证阶段，每做一次梯度下降，都要在测试集上评估一次\n",
    "        model.eval()\n",
    "        test_correct = 0\n",
    "        test_total = 0\n",
    "        \n",
    "        with torch.no_grad(): # 关键：禁用梯度计算\n",
    "            for inputs, labels in test_loader:\n",
    "                inputs, labels = inputs.to(device), labels.to(device)\n",
    "                outputs = model(inputs)\n",
    "                _, preds = torch.max(outputs, 1)\n",
    "                # torch.max返回(最大值, 最大值的索引)\n",
    "                # preds就是预测的类别索引，形状: [batch_size]（测试集的batchsize）\n",
    "                test_correct += (preds == labels).sum().item()\n",
    "                # (preds == labels) 返回布尔张量，True表示预测正确\n",
    "                # .sum() 计算True的数量，.item() 转换为Python数字\n",
    "                test_total += labels.size(0)\n",
    "\n",
    "        test_acc = test_correct / test_total\n",
    "        test_accuracies.append(test_acc)\n",
    "\n",
    "        # 输出进度\n",
    "        print(f'Epoch {epoch+1:2d}/{num_epochs} | '\n",
    "              f'Loss: {avg_loss:.4f} | '\n",
    "              f'Train Acc: {train_acc:.4f} | '\n",
    "              f'Test Acc: {test_acc:.4f}')\n",
    "\n",
    "        # 保存最佳模型\n",
    "        if test_acc > best_test_acc:\n",
    "            best_test_acc = test_acc\n",
    "            torch.save(model.state_dict(), \"best_model.pth\")\n",
    "            print(f\"↑ 保存最佳模型，准确率: {test_acc:.4f}\")\n",
    "\n",
    "    return train_losses, test_accuracies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4ef3a9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 开始训练\n",
    "# 定义损失函数和优化器\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "print(\"开始训练模型...\")\n",
    "print(\"=\" * 60)\n",
    "\n",
    "# 训练模型\n",
    "train_losses, test_accuracies = train_model(\n",
    "    model, train_loader, test_loader, criterion, optimizer, num_epochs=100\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68ba9daf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试模型预测\n",
    "def predict_image(model, image_path, transform, class_names):\n",
    "    \"\"\"预测单张图片\"\"\"\n",
    "    model.eval()\n",
    "\n",
    "    # 加载图片\n",
    "    image = Image.open(image_path).convert(\"RGB\")\n",
    "    image_tensor = transform(image).unsqueeze(0).to(device)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        outputs = model(image_tensor)\n",
    "        probabilities = torch.softmax(outputs, dim=1)\n",
    "        confidence, predicted = torch.max(probabilities, 1)\n",
    "\n",
    "        predicted_class = class_names[predicted.item()]\n",
    "        confidence_testue = confidence.item()\n",
    "\n",
    "    return predicted_class, confidence_testue\n",
    "\n",
    "\n",
    "# 测试预测功能\n",
    "def test_prediction():\n",
    "    # 从验证集中随机选择一张图片测试\n",
    "    import random\n",
    "\n",
    "    random_index = random.randint(0, len(test_dataset) - 1)\n",
    "    image_path, true_label = test_dataset.samples[random_index]\n",
    "    true_class = test_dataset.classes[true_label]\n",
    "\n",
    "    prediction, confidence = predict_image(\n",
    "        model, image_path, simple_test_transform, test_dataset.classes\n",
    "    )\n",
    "\n",
    "    # 显示图片和预测结果\n",
    "    image = Image.open(image_path).convert(\"RGB\")\n",
    "    plt.figure(figsize=(6, 6))\n",
    "    plt.imshow(image)\n",
    "    plt.title(f\"真实: {true_class}\\n预测: {prediction}\\n置信度: {confidence:.2f}\")\n",
    "    plt.axis(\"off\")\n",
    "    plt.show()\n",
    "\n",
    "    print(f\"真实类别: {true_class}\")\n",
    "    print(f\"预测类别: {prediction}\")\n",
    "    print(f\"置信度: {confidence:.2f}\")\n",
    "\n",
    "\n",
    "# 测试预测\n",
    "test_prediction()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
